"""Frequently used metrics

Currently including:

- compute_PSNR
- compute_SSIM
- compute_asa
- compute_ue
- compute_seg_category_counts
- compute_br
- compute_pri_and_voi

As suggested by the team, we are also going to include in the near future:

- compute_rmse
- compute_confusion_matrix
- compute_miou
- compute_rs_accuracy
"""

import math
import copy

import numpy as np
import torch
import torch.nn.functional as F
from common.ml_utils.seg_eval.compute_PRI_VoI import compute_RI_and_VoI_helper


def _gaussian(window_size, sigma):
    """create a gaussian filter"""
    gauss = torch.Tensor(
        [
            math.exp(-((x - window_size // 2) ** 2) / float(2 * sigma ** 2))
            for x in range(window_size)
        ]
    )
    return gauss / gauss.sum()


def _create_window(window_size, channel):
    """create a 2D moving window"""
    _1D_window = _gaussian(window_size, 1.5).unsqueeze(1)
    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
    window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
    return window


def compute_SSIM_tensor(
    tensor_pred: torch.Tensor,
    tensor_ref: torch.Tensor,
    window_size: int = 11,
    reduction="mean",
) -> float:
    """Compute SSIM (Structural Similarity Index Measure) from given prediction to given reference.

    NOTE this module is developed for computing with torch.Tensor in GPU more quickly rather
    than transferring tensor in GPU back to CPU, which is time-consuming.

    Parameters
    ----------
    tensor_pred : torch.Tensor
        Predicted image

        !!! note

        The given tensor is expected to be

        - a tensor of 4 dimensions, $n_{batches} * n_{bands} * n_{rows} * n_{cols}$
        - `tensor_pred` and `tensor_ref` should have the same dimension
        - in dtype `torch.float32`
        - within the value range of 0.0 to 1.0
        - in order to reduce time cost, we are *NOT* checking against those expectations within this function.
           It is the caller's duty to make sure those expectations are met.

    tensor_ref : torch.Tensor
        Reference image, sharing the same expectation as `tensor_pred`.
    window_size : int, optional
        window size used in computation, by default 11.
    reduction : str, optional
        Specifies the reduction to apply to the output: 'none' | 'mean'.
        'none': average will be taken for each image in the batch, the output will have `n_batches` SSIMs,
        'mean': average will be taken across all elements, the output will be just one SSIM.

    Returns
    -------
    float
        computed SSIM.

    Raises
    -------
    ValueError
        when `tensor_pred` and `tensor_ref` are not with the same shape,
        or they are ndarrays not with 4 dimensions.
    """
    if tensor_pred.shape != tensor_ref.shape:
        raise ValueError("`tensor_pred` and `tensor_ref` should have the same shape.")

    if len(tensor_pred.shape) != 4:
        raise ValueError("`tensor_pred` and `tensor_ref` should be in 4-dimensions.")

    img1 = tensor_pred
    img2 = tensor_ref
    (_, channel, _, _) = img1.size()
    window = _create_window(window_size, channel)
    window = window.type_as(img1)

    mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel)
    mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel)
    mu1_sq = mu1.pow(2)
    mu2_sq = mu2.pow(2)
    mu1_mu2 = mu1 * mu2
    sigma1_sq = (
        F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq
    )
    sigma2_sq = (
        F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq
    )
    sigma12 = (
        F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel)
        - mu1_mu2
    )
    C1 = 0.01 ** 2
    C2 = 0.03 ** 2
    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / (
        (mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)
    )

    if reduction == "mean":
        return ssim_map.mean()
    elif reduction == "none":
        return ssim_map.mean([1, 2, 3])
    else:
        raise ValueError("Invalid `reduction` method, accepts `mean` | `none`")


def compute_SSIM(
    nda_pred: np.ndarray,
    nda_ref: np.ndarray,
    window_size: int = 11,
    size_average: bool = True,
) -> float:
    """Compute SSIM (Structural Similarity Index Measure) from given prediction to given reference.

    Parameters
    ----------
    nda_pred : np.ndarray
        Predicted image

        !!! note

        The given array is expected to be

        - a ndarray of 4 dimensions, $n_{batches} * n_{bands} * n_{rows} * n_{cols}$
        - `nda_pred` and `nda_ref` should have the same dimension
        - in dtype `np.float32`
        - within the value range of 0.0 to 1.0
        - in order to reduce time cost, we are *NOT* checking against those expectations within this function.
           It is the caller's duty to make sure those expectations are met.

    nda_ref : np.ndarray
        Reference image, sharing the same expectation as `nda_pred`.
    window_size : int, optional
        window size used in computation, by default 11.
    size_average : bool, optional
        whether to average across the batch, by default True.

    Returns
    -------
    float
        computed SSIM.

    Raises
    -------
    ValueError
        when `nda_pred` and `nda_ref` are not with the same shape, or they are ndarrays not with 4 dimensions.
    """
    if nda_pred.shape != nda_ref.shape:
        raise ValueError("`nda_pred` and `nda_ref` should have the same shape.")

    if len(nda_pred.shape) != 4:
        raise ValueError("`nda_pred` and `nda_ref` should be in 4-dimensions.")

    img1 = torch.Tensor(nda_pred)
    img2 = torch.Tensor(nda_ref)
    (_, channel, _, _) = img1.size()
    window = _create_window(window_size, channel)
    window = window.type_as(img1)

    mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel)
    mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel)
    mu1_sq = mu1.pow(2)
    mu2_sq = mu2.pow(2)
    mu1_mu2 = mu1 * mu2
    sigma1_sq = (
        F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq
    )
    sigma2_sq = (
        F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq
    )
    sigma12 = (
        F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel)
        - mu1_mu2
    )
    C1 = 0.01 ** 2
    C2 = 0.03 ** 2
    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / (
        (mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)
    )

    if size_average:
        return ssim_map.mean().item()
    else:
        return ssim_map.mean(1).mean(1).mean(1).numpy()


def compute_PSNR_tensor(
    tensor_pred: torch.Tensor, tensor_ref: torch.Tensor, reduction="mean"
) -> float:
    """Compute PSNR (Peak Signal Noise Ratio) from given prediction to given reference.

    NOTE this module is developed for computing with torch.Tensor in GPU more quickly rather
    than transferring tensor in GPU back to CPU, which is time-consuming.

    Parameters
    ----------
    tensor_pred : torch.Tensor
        Predicted image

        !!! note

        The given tensor is expected to be

        - a tensor of 4 dimensions, $n_{batches} * n_{bands} * n_{rows} * n_{cols}$
          or $n_{batches} * n_{rows} * n_{cols} * n_{bands}$
        - or a tensor of 3 dimensions, $n_{rows} * n_{cols} * n_{bands}$ or $n_{bands} * n_{rows} * n_{cols}$
        - or a tensor of 2 dimensions, $n_{rows} * n_{cols}$
        - `tensor_pred` and `tensor_ref` should have the same dimension
        - in dtype `torch.float32`
        - within the value range of 0.0 to 1.0
        - in order to reduce time cost, we are *NOT* checking against those expectations within this function.
           It is the caller's duty to make sure those expectations are met.

    tensor_ref : torch.Tensor
        Reference image, sharing the same expectation as `tensor_pred`.

    reduction : str, optional
        Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'.
        'none': the output will have `n_batches` PSNRs for each image in the batch,
        'mean': the average of PSNRs across all images in the batch, the output will be just one PSNR,
        'sum': the sum of PSNRs across all images in the batch, the output will be just one PSNR.

    Returns
    -------
    float
        computed PSNR.

    Raises
    -------
    ValueError
        when `tensor_pred` and `tensor_ref` are not with the same shape,
        or they are tensors with more than 4 dimensions.

    Notes
    -------
    We calculate PSNR using the following formula,
    from [wikipedia](https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio):
    $$
    psnr = 20 * log(\\frac{MAX}{\\sqrt{MSE}})
    $$

    Notice since we expect the input to be ndarray with dtype `np.float32`
    and within the value range 0.0 to 1.0, $MAX$ is 1.0.
    If there is no difference between the nda_pred and nda_ref, we assume the PSNR to be 100.0.

    Since rmse ranges [0, 1], psnr ranges [0, Inf]
    """
    if tensor_pred.shape != tensor_ref.shape:
        raise ValueError("`tensor_pred` and `tensor_ref` should have the same shape.")

    def _compute(tensor_pred, tensor_ref):
        tensor_diff = tensor_pred - tensor_ref
        rmse = torch.sqrt(torch.mean(tensor_diff ** 2))
        if rmse == 0:
            psnr = 100.0
        else:
            psnr = 20.0 * torch.log10(1.0 / rmse)
        return psnr

    if len(tensor_pred.shape) > 4:
        raise ValueError("Expect input tensor to be of 4 dimension or less.")
    elif len(tensor_pred.shape) == 4:
        psnrs = []
        for i in range(tensor_pred.shape[0]):
            psnrs.append(_compute(tensor_pred[i, ...], tensor_ref[i, ...]).item())
        psnrs = torch.FloatTensor(psnrs).to(device=tensor_pred.device)
        if reduction == "none":
            return psnrs
        elif reduction == "mean":
            return torch.mean(psnrs)
        elif reduction == "sum":
            return torch.sum(psnrs)
        else:
            raise ValueError(
                "Invalid `reduction` method, accepts `none` | `mean` | `sum`."
            )
    elif len(tensor_pred.shape) <= 3:
        return _compute(tensor_pred, tensor_ref)


def compute_PSNR(nda_pred: np.ndarray, nda_ref: np.ndarray) -> float:
    """Compute PSNR (Peak Signal Noise Ratio) from given prediction to given reference.

    Parameters
    ----------
    nda_pred : np.ndarray
        Predicted image

        !!! note

        The given array is expected to be

        - a ndarray of 4 dimensions, $n_{batches} * n_{bands} * n_{rows} * n_{cols}$
          or $n_{batches} * n_{rows} * n_{cols} * n_{bands}$
        - or a ndarray of 3 dimensions, $n_{rows} * n_{cols} * n_{bands}$ or $n_{bands} * n_{rows} * n_{cols}$
        - or a ndarray of 2 dimensions, $n_{rows} * n_{cols}$
        - `nda_pred` and `nda_ref` should have the same dimension
        - in dtype `np.float32`
        - within the value range of 0.0 to 1.0
        - in order to reduce time cost, we are *NOT* checking against those expectations within this function.
           It is the caller's duty to make sure those expectations are met.

    nda_ref : np.ndarray
        Reference image, sharing the same expectation as `nda_pred`.

    Returns
    -------
    float
        computed PSNR.

    Raises
    -------
    ValueError
        when `nda_pred` and `nda_ref` are not with the same shape, or they are ndarrays with more than 4 dimensions.

    Notes
    -------
    We calculate PSNR using the following formula,
    from [wikipedia](https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio):
    $$
    psnr = 20 * log(\\frac{MAX}{\\sqrt{MSE}})
    $$

    Notice since we expect the input to be ndarray with dtype `np.float32`
    and within the value range 0.0 to 1.0, $MAX$ is 1.0.
    If there is no difference between the nda_pred and nda_ref, we assume the PSNR to be 100.0.

    Since rmse ranges [0, 1], psnr ranges [0, Inf]
    """
    if nda_pred.shape != nda_ref.shape:
        raise ValueError("`nda_pred` and `nda_ref` should have the same shape.")

    def _compute(nda_pred, nda_ref):
        nda_diff = nda_pred - nda_ref
        rmse = math.sqrt(np.mean(nda_diff ** 2))
        if rmse == 0:
            psnr = 100.0
        else:
            psnr = 20.0 * math.log10(1.0 / rmse)
        return psnr

    if len(nda_pred.shape) > 4:
        raise ValueError("Expect input array to be of 4 dimension or less.")
    elif len(nda_pred.shape) == 4:
        psnr = 0
        for i in range(nda_pred.shape[0]):
            psnr += _compute(nda_pred[i, ...], nda_ref[i, ...])
        return psnr / nda_pred.shape[0]
    elif len(nda_pred.shape) <= 3:
        return _compute(nda_pred, nda_ref)


def miou_multiclass(arr1: np.ndarray, arr2: np.ndarray, valid_class):
    """
    Calculate the MIOU between two ndarrays. both arrays should set 0 as nodata class.
    the shape of arrays should be the same.
    only class in valid_class will be used to calculate miou.
    0 in valid_class will always be used as nodata.
    other class number not listed in valid_class will be regarded as 0 (nodata) and nodata
    will not be used to calculate miou.

    return:
        float # the MIOU value of two arrays. the shape will be a float number in [0, 1]

    """

    valid_class = set(valid_class)

    tmp_arr1 = copy.deepcopy(arr1)
    tmp_arr2 = copy.deepcopy(arr2)

    tmp_arr1 = np.squeeze(tmp_arr1)
    tmp_arr2 = np.squeeze(tmp_arr2)

    if tmp_arr1.shape != tmp_arr2.shape:
        raise ValueError(
            "array1 and array2 should have same shape, but have {} and {}".format(
                tmp_arr1.shape, tmp_arr2.shape
            )
        )
    tmp_arr1 = tmp_arr1.reshape(-1)
    tmp_arr2 = tmp_arr2.reshape(-1)

    if not np.issubdtype(tmp_arr1.dtype, np.integer):
        raise ValueError("array1 dtype is not integer")

    if not np.issubdtype(tmp_arr2.dtype, np.integer):
        raise ValueError("array2 dtype is not integer")

    tmp_arr1[~np.isin(tmp_arr1, list(valid_class))] = 0
    tmp_arr2[~np.isin(tmp_arr2, list(valid_class))] = 0
    valid_mask = np.logical_and(tmp_arr1 != 0, tmp_arr2 != 0)

    iou_result = 0.0
    for class_idx in valid_class:
        mask_arr1 = np.logical_and(tmp_arr1 == class_idx, valid_mask)
        mask_arr2 = np.logical_and(tmp_arr2 == class_idx, valid_mask)
        arr_inter = np.sum(np.logical_and(mask_arr1, mask_arr2))
        arr_union = np.sum(np.logical_or(mask_arr1, mask_arr2))
        # Handle the corner case when arr_union is zero.
        iou_result += 1.0 if arr_union == 0 else arr_inter / arr_union
    return iou_result / len(valid_class)


def preprocess_nodata(prediction, ground_truth):
    valid_mask = np.logical_and(prediction != 0, ground_truth != 0)
    pos_predict = np.logical_and(prediction == 2, valid_mask)
    neg_predict = np.logical_and(prediction == 1, valid_mask)

    pos_GT = np.logical_and(ground_truth == 2, valid_mask)
    neg_GT = np.logical_and(ground_truth == 1, valid_mask)

    return pos_predict, neg_predict, pos_GT, neg_GT


def general_accuracy_binary(prediction: np.ndarray, ground_truth: np.ndarray):
    pos_pred, neg_pred, pos_gt, neg_gt = preprocess_nodata(prediction, ground_truth)

    num_TP = np.sum(np.logical_and(pos_pred, pos_gt))
    num_FP = np.sum(np.logical_and(pos_pred, neg_gt))
    num_TN = np.sum(np.logical_and(neg_pred, neg_gt))
    num_FN = np.sum(np.logical_and(neg_pred, pos_gt))
    return (num_TP + num_TN) / (num_TP + num_TN + num_FP + num_FN)


def recall_binary(prediction: np.ndarray, ground_truth: np.ndarray):
    pos_pred, neg_pred, pos_gt, _ = preprocess_nodata(prediction, ground_truth)

    num_TP = np.sum(np.logical_and(pos_pred, pos_gt))
    num_FN = np.sum(np.logical_and(neg_pred, pos_gt))
    return num_TP / (num_TP + num_FN)


def precision_binary(prediction: np.ndarray, ground_truth: np.ndarray):
    pos_pred, _, pos_gt, neg_gt = preprocess_nodata(prediction, ground_truth)

    num_TP = np.sum(np.logical_and(pos_pred, pos_gt))
    num_FP = np.sum(np.logical_and(pos_pred, neg_gt))
    return num_TP / (num_TP + num_FP)


def f1_score_binary(prediction: np.ndarray, ground_truth: np.ndarray):
    precision_val = precision_binary(prediction, ground_truth)
    recall_val = recall_binary(prediction, ground_truth)
    return (2 * precision_val * recall_val) / (precision_val + recall_val)


def miou_binary(prediction: np.ndarray, ground_truth: np.ndarray):
    pos_pred, neg_pred, pos_gt, neg_gt = preprocess_nodata(prediction, ground_truth)

    inter_true = np.sum(np.logical_and(pos_pred, pos_gt))
    union_true = np.sum(np.logical_or(pos_pred, pos_gt))

    inter_false = np.sum(np.logical_and(neg_pred, neg_gt))
    union_false = np.sum(np.logical_or(neg_pred, neg_gt))

    miou = ((inter_true / union_true) + (inter_false / union_false)) * 0.5
    return miou


def compute_asa(label_map, gt_labels_map):
    """Computer ASA (Achievable Segmentation Accuracy)

    Parameters
    ----------
    label_map : ndarray
        the segmentation resultant label map
    gt_labels_map : ndarray
        the ground truth map

    Returns
    -------
    asa : float
        the ASA value
    """
    label_list = label_map.flatten().tolist()
    gtseg_list = gt_labels_map.flatten().tolist()

    import EvalSPModule

    asa = EvalSPModule.computeASA(label_list, gtseg_list, 0)
    return asa


def compute_ue(label_map, gt_labels_map):
    """Computer UE (Bergh Underseg Error, 1 - ASA)

    Parameters
    ----------
    label_map : ndarray
        the segmentation resultant label map
    gt_labels_map : ndarray
        the ground truth map

    Returns
    -------
    ue : float
        the UE value
    """
    label_list = label_map.flatten().tolist()
    gtseg_list = gt_labels_map.flatten().tolist()

    import EvalSPModule

    asa = EvalSPModule.computeASA(
        label_list, gtseg_list, 0
    )  # the last parameter is for return the label map, default is 0
    ue = 1 - asa
    return ue


def compute_seg_category_counts(label_map, gt_labels_map, perc):
    """Computer Segmentation counts in three categories

    Parameters
    ----------
    label_map : ndarray
        the segmentation resultant label map
    gt_labels_map : ndarray
        the ground truth map
    perc : float [0 <= perc <= 1]
        the overlap percentage perc

    Returns
    -------
    nr_underseg : int
        the number of under-segmentation under the overlap percentage perc

    nr_one_to_one : int
        the number of one-to-one-segmentation under the overlap percentage perc

    nr_overseg : int
        the number of over-segmentation under the overlap percentage perc

    nr_gt : int
        the number of ground truth segments

    Raises
    ------
    ValueError
        perc must be 0 <= perc <= 1
    """

    label_list = label_map.flatten().tolist()
    gtseg_list = gt_labels_map.flatten().tolist()

    if perc < 0 or perc > 1:
        raise ValueError("the input perc must be within 0.0 and 1.0.")

    import EvalSPModule

    nr_underseg, nr_one_to_one, nr_overseg, nr_gt = EvalSPModule.computeSegCatCounts(
        label_list, gtseg_list, perc, 0
    )  # the last parameter is for return the label map, default is 0
    return nr_underseg, nr_one_to_one, nr_overseg, nr_gt


def compute_br(label_map, gt_labels_map, search_radius):
    """Computer BR (Boundary Recall)

    Parameters
    ----------
    label_map : ndarray
        the segmentation resultant label map
    gt_labels_map : ndarray
        the ground truth map
    search_radius : int
        the searching radius of a successful recall

    Returns
    -------
    br_gt : float [0<= br_gt <= 1]
        BR (boundary recall with regard to GT)
    br_sp : float [0<= br_gt <= 1]
        BR (boundary recall with regard to segmentation result)
    average_error_pixel_gt: float
        Average Error Pixel of segmentation result with regard to GT
    average_error_pixel_sp: float
        Average Error Pixel of segmentation result with regard to segmentation result

    Raises
    ------
    ValueError
        search_radius must be larger than 0
    """

    h, w = label_map.shape

    label_list = label_map.flatten().tolist()
    gtseg_list = gt_labels_map.flatten().tolist()

    if search_radius < 0:
        raise ValueError("the input search_radius must be larger than 0.")

    import EvalSPModule

    (
        br_gt,
        br_sp,
        average_error_pixel_gt,
        average_error_pixel_sp,
    ) = EvalSPModule.computeBR(label_list, gtseg_list, h, w, search_radius)
    return br_gt, br_sp, average_error_pixel_gt, average_error_pixel_sp


def compute_pri_and_voi(label_map, gt_labels_map):
    """Compute PRI (Probabilistic Rand Index) and VoI (Variation of Information) for segmentation performance evaluation.

    These metrics are suitable for the panoptic segmentation scenario but the ground truth classes are not clear.
    Therefore, evaluation of segmentation performance has to be achieved by purely comparing
    the segmentation consistency between segmentation results and GT.

    Reference:
    Li, Z., Wu, X. M., & Chang, S. F. (2012, June).
    Segmentation using superpixels: A bipartite graph partitioning approach.
    In 2012 IEEE conference on computer vision and pattern recognition (pp. 789-796). IEEE.
    https://github.com/ColumbiaDVMM/Segmentation-Using-Superpixels/blob/master/evals/compare_segmentations.m
    https://github.com/ColumbiaDVMM/Segmentation-Using-Superpixels/blob/master/evals/eval_segmentation.m

    The definitions of PRI and VoI can be found in:
    Unnikrishnan, R., Pantofaru, C., & Hebert, M. (2007). Toward objective evaluation of image segmentation algorithms.
    IEEE transactions on pattern analysis and machine intelligence, 29(6), 929-944.

    Parameters
    ----------
    label_map : ndarray
        segmentation label map
    gt_labels_map : list or ndarray
        ground truth label map(s)

    Returns
    -------
    PRI : float
        PRI result [0, 1]
    VoI : float
        VoI result >= 0
    """
    (tx, ty) = label_map.shape
    num1 = np.max(label_map)
    if isinstance(gt_labels_map, list):
        PRI = 0.0
        VoI = 0.0
        for tmp_gt in gt_labels_map:
            num2 = np.max(tmp_gt)
            curRI, curVoI = compute_RI_and_VoI_helper(
                label_map, tmp_gt, tx, ty, num1, num2
            )
            PRI += curRI
            VoI += curVoI
        PRI /= len(gt_labels_map)
        VoI /= len(gt_labels_map)
    else:
        num2 = np.max(gt_labels_map)
        PRI, VoI = compute_RI_and_VoI_helper(
            label_map, gt_labels_map, tx, ty, num1, num2
        )

    return PRI, VoI


if __name__ == "__main__":
    from PIL import Image

    pred = np.array(Image.open("./examples/norain.png"))
    gt = np.array(Image.open("./examples/rain.png"))

    psnr = compute_PSNR(pred / 255.0, gt / 255.0)
    print("PSNR: {}".format(psnr))

    pred = torch.Tensor(pred).permute(2, 0, 1).unsqueeze(0).repeat(16, 1, 1, 1)
    gt = torch.Tensor(gt).permute(2, 0, 1).unsqueeze(0).repeat(16, 1, 1, 1)

    # expect inputs in ndarray and clamp to 0 ~ 1.0
    pred = pred.data.cpu().numpy().astype(np.float32) / 255
    gt = gt.data.cpu().numpy().astype(np.float32) / 255
    psnr = compute_PSNR(pred, gt)
    print("PSNR: {}".format(psnr))

    ssim = compute_SSIM(pred, gt, size_average=False)
    print("SSIM: {}".format(ssim))

    # test for miou
    test1 = np.array([1, 2, 3, 1, 2, 3, 2, 1], dtype=np.float32)
    test2 = np.array([1, 2, 3, 1, 2, 3, 2, 1])
    test3 = np.array([1, 1, 3, 1, 1, 1, 1, 1])
    print(test2.dtype, test3.dtype)
    print(miou_binary(test2, test3, [1, 3]))  # (3/7 + 1/2)/2
